import numpy as np
from typing import Literal
from ultralytics.models.yolo import YOLO
from pathlib import Path
PATH_PARENT = Path(__file__).parent
PATH_MODELS = PATH_PARENT / "../yolo_models"

class YoloDetection:
    def __init__(
        self,
        model_name: Literal[
            '618_surface_detect_m', 'surface_detect_387epoch'
        ] = 'surface_detect_387epoch',
        engine: Literal['torch', 'onnx', 'tensorrt'] = 'torch',
    ):
        self.engine = engine
        model_path = PATH_MODELS / f"{model_name}.pt"
        if engine != 'torch':
            model_path = model_path.with_suffix('.onnx' if engine == 'onnx' else '.engine')
            if not model_path.exists():
                print(f"Can't find {model_path}, automatically convert")
                yolo_pt = YOLO(PATH_MODELS / f"{model_name}.pt")
                yolo_pt.export(format=engine, half=True, imgsz=[480, 640])
                print(f"Convert done, saved to {model_path}")

        print(f"Load model from {model_path} with engine {engine}")
        self.model = YOLO(model_path)
    
    def detect(self, img: np.ndarray):
        results = self.model(img, verbose=False, conf=0.6, iou=0.3)[0]
        boxes = results.boxes
        render_img = results.plot()
        return boxes, render_img

if __name__ == '__main__':
    import cv2
    from manipulation.scripts.detection.cameras.d435_camera import D435Camera

    yolo = YoloDetection()
    camera = D435Camera(rgb_width=1920, rgb_height=1080, fps=15, only_rgb=True)
    while True:
        rgb, _ = camera.get_frame()
        boxes, render_img = yolo.detect(rgb)
        cv2.imshow("detect", render_img)
        cv2.waitKey(1)
